#!/usr/bin/env python
# coding: utf-8

# # scvelo basics
# 
# from https://scvelo.readthedocs.io/en/stable/VelocityBasics/
# I used this to create this conda env: 
# `mamba create -n scvelo scvelo python-igraph jupyterlab "pandas=1.5.3" "numpy=1.23.5" "matplotlib=3.3.5"`

# if in rstudio: reticulate::use_condaenv('scvelo-colorbar')

import scvelo as scv
scv.logging.print_version()

scv.settings.verbosity = 3  # show errors(0), warnings(1), info(2), hints(3)
scv.settings.presenter_view = True  # set max width size for presenter view
scv.set_figure_params('scanpy')  # for beautified visualization

# ## load the data (will create a data dir under work dir)

adata = scv.datasets.pancreas()
adata = scv.read_loom('/home/gjsx/append-ssd/nextflowing/scrna-sle-perez2022-p2.2-3/cellranger/count/CLUES1_POOL02_2/velocyto/CLUES1_POOL02_2.loom')

adata

scv.pl.proportions(adata)

# ## preprocess the data
# First order is needed for deterministic velocity estimation,
# while stochastic estimation also requires second order moments.
scv.pp.filter_and_normalize(adata, min_shared_counts=20, n_top_genes=2000)
scv.pp.moments(adata, n_pcs=30, n_neighbors=30)

# ## estimate RNA velocity

scv.tl.velocity(adata)

# this step will fail if numpy is not degraded to 1.23.5
scv.tl.velocity_graph(adata)

# ## Project the velocity

# this step will fail if pandas is not degrade
scv.pl.velocity_embedding_stream(adata, basis='pca')

scv.pl.velocity_embedding(adata, arrow_length=3, arrow_size=2, dpi=120)

# ## Inteprete the velocity

# downgrade matplotlib to 3.3.4 to avoid error in matplotlib
# or, add colorbar=False to pl command to disable all continuous color legends
scv.pl.velocity(adata, ['Cpe',  'Gnao1', 'Ins2', 'Adk'], ncols=2)

scv.pl.scatter(adata, 'Cpe', color=['clusters', 'velocity'],
               add_outline='Ngn3 high EP, Pre-endocrine, Beta')

# ## Identify important genes

scv.tl.rank_velocity_genes(adata, groupby='clusters', min_corr=.3)

df = scv.core.get_df(adata.uns['rank_velocity_genes']['names'])
df

kwargs = dict(frameon=False, size=10, linewidth=1.5)

scv.pl.scatter(adata, df['Ngn3 high EP'][:5],
ylabel='Ngn3 high EP',add_outline='Ngn3 high EP', **kwargs)

scv.pl.scatter(adata, df['Pre-endocrine'][:5], ylabel='Pre-endocrine',
add_outline='Pre-endocrine', **kwargs)

# ## Velocities in cycling progenitors

scv.tl.score_genes_cell_cycle(adata)
# can see ductal cells are actively dividing
scv.pl.scatter(adata, color_gradients=['S_score', 'G2M_score'], smooth=True, perc=[5, 95])

s_genes = scv.get_df(adata[:, s_genes], 'spearmans_score', sort_values=True).index
g2m_genes = scv.get_df(adata[:, g2m_genes], 'spearmans_score', sort_values=True).index

## first 2 s_genes & 2 g2m genes
scv.pl.scatter(adata, list(s_genes[:2]) + list(g2m_genes[:2]),
ncols=2, add_outline='Ductal', hspace=.5)

# if set add_outline=True, all cell will have outline
# Top2a upregulated in S phase, downregulated in G2M phase
scv.pl.velocity(adata, ['Top2a'], add_outline='Ductal')

# ## Speed and coherence --------------
## length of velocity indicates differentiation speed
## coherence of the vector field (i.e., how a velocity vector correlates with its neighboring velocities)
## provides a measure of confidence
scv.tl.velocity_confidence(adata)
keys = 'velocity_length', 'velocity_confidence'

# c stands for color, cmap stands for color_map
scv.pl.scatter(adata, c=keys, cmap='coolwarm', perc=[5, 95])

# high confidence show a determined differentitation direction

# see obs table, group by clusters, extract keys (length & confidence) column
# calc means, transpose
df = adata.obs.groupby('clusters')[keys].mean().T

df.style.background_gradient(cmap='coolwarm', axis=1)
# a heatmap like table, not displayed in rstudio

# ## Velocity graph & pseudotime ------------

# show possible transition path between cells
# higher threshold eliminate low probability transition
scv.pl.velocity_graph(adata, threshold=.2)
# connecting line is dislocated here...
# downgrade matplotlib fix it!

# draw descendents/anscestors coming from a specified cell
x, y = scv.utils.get_cell_transitions(adata, basis='umap', starting_cell='AAACCTGAGCCTTGAT')
ax = scv.pl.velocity_graph(adata, c='lightgrey', edge_width=.05, show=False)
scv.pl.scatter(adata, x=x, y=y, s=120, c='ascending', cmap='gnuplot', ax=ax)
# a very small plot?
# downgrade matplotlib!

# infer pseudotime from all vector field
scv.tl.velocity_pseudotime(adata)
scv.pl.scatter(adata, color='velocity_pseudotime', cmap='gnuplot')

# ## PAGA velocity graph -----------

# requires igraph
# !mamba install -n scvelo python-igraph

# PAGA graph abstraction has benchmarked as top-performing method for trajectory inference
# velocity-inferred directionality between clusters
scv.tl.paga(adata, groups='clusters')
df = scv.get_df(adata, 'paga/transitions_confidence', precision=2).T
df
df.style.background_gradient(cmap='Blues').format('{:.2g}')

scv.pl.paga(adata, basis='umap', size=50, alpha=.1,
            min_edge_width=2, node_size_scale=1.5)

